from ultralytics import YOLO
from custom_trainer import CustomTrainer

YOLO("/home/cui_4090/ultralytics/ultralytics/nn/modules/yolo11n.pt").train(
    trainer=CustomTrainer,
    data="/home/cui_4090/ultralytics/dataset/data.yaml",  # 确保 nc=1, names: ['panel']
    imgsz=1024,
    epochs=300,
    batch=16,           # 按显存调
    workers=8,
    device=0            # GPU id
)
